import os
import csv
from zkl_aiutils_datasets import load_dataset
from zkl_pyutils_fsspec import FsLike, resolve_fs

class LabelEmbeddingReader:
    def __init__(self, dataset_path):
        self.dataset_path = dataset_path
        self.dataset = self.dataset_read_init(dataset_path)
        self.key_to_id = self.csv_init(dataset_path)

    def dataset_read_init(self, dataset_path):
        fs = resolve_fs(os.path.join(dataset_path, "dataset"))
        dataset = load_dataset(fs)
        return dataset

    def csv_init(self, dataset_path):
        key_to_id = {}
        csv_path = os.path.join(dataset_path, "label_keys.csv")
        if not os.path.exists(csv_path):
            raise FileNotFoundError(f"Could not find label_keys.csv at {csv_path}")
        with open(csv_path, 'r') as f:
            reader = csv.reader(f)
            next(reader)  # Skip header
            for i, row in enumerate(reader):
                key_to_id[row[0]] = i
        return key_to_id

    def get_labels_embeddings_batch(self, keys):
        embeddings = []
        for key in keys:
            try:
                id = self.key_to_id[key]
                embeddings.append(self.dataset[id][0][0])
            except KeyError:
                embeddings.append(None)
        return embeddings

    def get_label_embedding(self, key):
        embeddings = self.get_labels_embeddings_batch([key])
        return embeddings[0] if embeddings else None

    def iterate_embeddings(self):
        """Iterates over all keys and their corresponding embeddings."""
        for key, index in self.key_to_id.items():
            yield key, self.dataset[index][0]
    
    def count_embeddings(self):
        return len(self.key_to_id)

    
    def close(self):
        # self.dataset.close()
        pass

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()

    def __del__(self):
        self.close()

if __name__ == '__main__':
    project_dir_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
    dataset_path = os.path.join(project_dir_path, "embeddings_local/ukbiobank/v4-labels_embeddings")
    print(f"dataset_path: {dataset_path}")
    
    print(f"Attempting to read from: {dataset_path}")
    
    try:
        reader = LabelEmbeddingReader(dataset_path=dataset_path)
        
        keys_to_test = ["p48_i0", "p51_i2", "p23480_i1", "non_existent_field"]
        print(f"\nTesting with keys: {keys_to_test}")
        embeddings = reader.get_labels_embeddings_batch(keys_to_test)
        for key, emb in zip(keys_to_test, embeddings):
            status = "Found" if emb is not None else "Not Found"
            # print(f"  - Embedding for key {key}: {status}")
            print(f"  - Embedding for key {key}: {emb}")

        key_to_test = "p48_i0"
        print(f"\nTesting with single key: {key_to_test}")
        embedding = reader.get_label_embedding(key_to_test)
        status = "Found" if embedding is not None else "Not Found"
        print(f"  - Embedding for key {key_to_test}: {status}")

    except FileNotFoundError as e:
        print(f"\nError: {e}")
        print("Please ensure you have run the make_labels.py and processing_labels.py scripts first.")